from torch.utils.data import DataLoader
from torch import optim
import torch
import os
from torch import nn
from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.Resize(size=58),
    transforms.CenterCrop(size=58),
    transforms.ToTensor()
])


def accury(predictions, labels):
    pred = torch.max(predictions.data, 1)[1]
    rights = pred.eq(labels.data.view_as(pred)).sum()
    return rights, len(labels)


class GenderModule(nn.Module):
    def __init__(self):
        super(GenderModule, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=3,
                      kernel_size=3,
                      stride=1,
                      padding=1,
                      out_channels=32), nn.ReLU(), nn.MaxPool2d(kernel_size=2))
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=32,
                      kernel_size=3,
                      stride=1,
                      padding=1,
                      out_channels=128), nn.ReLU(),
            nn.MaxPool2d(kernel_size=2))
        self.out = nn.Linear(in_features=128 * 7 * 7, out_features=2)

    def forward(self, x):
        print(x)
        x = self.conv1(x)
        print(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        x = self.out(x)
        return x


train_image = datasets.ImageFolder(os.path.join(os.getcwd(), "Training"),
                                   transform=transform)
valid_image = datasets.ImageFolder(os.path.join(os.getcwd(), "Validation"),
                                   transform=transform)
training_male_dataset = DataLoader(dataset=train_image,
                                   batch_size=8,
                                   shuffle=True)

if __name__ == '__main__':
    net = GenderModule()
    lost_func = nn.CrossEntropyLoss()
    op = optim.Adam(params=net.parameters(), lr=0.001)

    for i in range(1000):
        train_rights = []

        for inx, (data, label) in enumerate(training_male_dataset):
            print(data.shape)
            output = net(data)
            loss = lost_func(output, label)
            op.zero_grad()
            loss.backward()
            op.step()
            right = accury(output, labels=label)
            train_rights.append(right)

            if i % 100 == 0:
                print('training {}' % i)
